"""React 智能体的工具管理器"""
from typing import Dict, List, Any
from tools import CalculatorTool, BaiduSearchTool, BingSearchTool
from tools import ToolResult
from langchain_mcp_adapters.client import MultiServerMCPClient
from config import Config


class ToolManager:
    """React 智能体的工具管理器"""
    
    def __init__(self):
        self.tools: Dict[str, Any] = {}
        self._initialize_tools()
        self.user_tools: Dict[str, List[str]] = {}
        self.client = None

    def _initialize_tools(self):
        """初始化所有可用工具 """
        tools = [
            BaiduSearchTool(),
            # BingSearchTool(),
            CalculatorTool(),
        ]
        self.tools = {tool.name: tool for tool in tools}

    async def load_mcp_tools(self, config: Dict, override: bool = False) -> List[str]:
        self.client = MultiServerMCPClient(config)
        tools = await self.client.get_tools()
        if override:
            self.tools = {tool.name: tool for tool in tools}
        else:
            self.tools.update({tool.name: tool for tool in tools})
        return list(self.tools.keys())

    async def load_session_tools(self, session_id: str, tools: List[str] = None):
        """加载用户的工具"""
        if self.client is None:
            config = {
                "system_mcp": {
                    "url": Config.MCP_URI,
                    "transport": "sse",
                    "sse_read_timeout": 10 * 60,
                },
                "gaode_mcp": {
                    "url": Config.GAODE_MCP_URI,
                    "transport": "sse",
                    "sse_read_timeout": 10 * 60,
                },
            }
            await self.load_mcp_tools(config, override=False)

        if tools:
            self.user_tools[session_id] = [tool for tool in tools if tool in self.tools]
        else:
            self.user_tools[session_id] = list(self.tools.keys())
        return self.user_tools[session_id]

    def get_tool_names(self, session_id: str = None) -> List[str]:
        if isinstance(session_id, str) and session_id in self.user_tools:
            return self.user_tools[session_id]
        else:
            return list(self.tools.keys())

    async def execute_tool(self, tool_name: str, tool_input: Dict, **kwargs) -> ToolResult:

        if tool_name not in self.tools:
            return ToolResult(success=False, data=None,
                error=f"未找到工具'{tool_name}'. 可用工具列表: {', '.join(self.get_tool_names())}"
            )
        try:
            tool = self.tools.get(tool_name)
            tool_input.update(kwargs)
            result = await tool.ainvoke(tool_input)

            if isinstance(result, ToolResult):
                return result
            else:
                return ToolResult(success=True,data=result,error=None)

        except Exception as e:
            return ToolResult(success=False, data=None, error=f"工具调用失败：{str(e)}")

    def format_tools_for_prompt(self, session_id: str = None) -> List[Dict]:
        """将工具信息格式化，以便纳入提示词中"""
        tool_descriptions = []

        if isinstance(session_id, str) and session_id in self.user_tools:
            user_tools = self.user_tools[session_id]
        else:
            user_tools = self.tools.keys()

        for name in user_tools:
            tool = self.tools.get(name)
            tool_descriptions.append({
                'tool.name': name,
                'tool.description': str(tool.description),
                'tool.args_schema': str(tool.args_schema)
            })

        return tool_descriptions